from FlagEmbedding import FlagReranker

def rerank(query, chunks):
    reranker_path = "../model/OpenScholar_Reranker"
    reranker = FlagReranker(reranker_path, use_fp16=True)
    # scores = reranker.compute_score([[query, c["snippets"][0]] for c in chunks], batch_size=100)
    scores = reranker.compute_score([[query, c] for c in chunks], batch_size=100)
    if type(scores) is float:
        result_dic = {0: scores}
    else:
        result_dic = {p_id: score for p_id, score in enumerate(scores)}
    p_ids = sorted(result_dic.items(), key=lambda x: x[1], reverse=True)
    new_orders = []
    indexs = []
    for i, p_id in enumerate(p_ids):
        new_orders.append(chunks[p_id[0]])
        indexs.append(p_id[0])
    return new_orders,indexs
    
def rm_rerank(query, search_results):
    reranker_path = "../model/OpenScholar_Reranker"
    reranker = FlagReranker(reranker_path, use_fp16=True)
    scores = reranker.compute_score([[query, r["entity"]["chunk_text"]] for r in search_results], batch_size=100)
    if type(scores) is float:
        result_dic = {0: scores}
    else:
        result_dic = {p_id: score for p_id, score in enumerate(scores)}
    p_ids = sorted(result_dic.items(), key=lambda x: x[1], reverse=True)
    new_orders = []
    indexs = []
    for i, p_id in enumerate(p_ids):
        new_orders.append(search_results[p_id[0]])
        indexs.append(p_id[0])
    return new_orders